package controller

import (
	"context"
	"encoding/json"
	"fmt"
	"github.com/gin-gonic/gin"
	"github.com/tmc/langchaingo/llms"
	"net/http"
	"test/ollama/model"
	"test/ollama/util"
)

// 只为controller调用
type Chat struct {
}

func (chat Chat) DoChat(c *gin.Context) {
	//进行数据绑定
	var body model.Chat
	err := c.ShouldBindJSON(&body)

	if err != nil {
		fmt.Println("绑定传入json数据信息异常", err)
		c.JSON(http.StatusBadRequest, gin.H{
			"msg": err.Error(),
		})
		return
	}

	prompt := util.Prompt()
	data := map[string]any{
		"text": body.Text,
	}
	//格式化大模型数据
	msg, _ := prompt.FormatMessages(data)
	//设置大模型数据信息
	contents := []llms.MessageContent{
		llms.TextParts(msg[0].GetType(), msg[0].GetContent()),
		llms.TextParts(msg[1].GetType(), msg[1].GetContent()),
		llms.TextParts(msg[2].GetType(), msg[2].GetContent()),
	}
	//oll := util.Ollama(c, "llama3")
	oll := util.Ollama(c, "qwen2")
	content, err := oll.GenerateContent(context.Background(), contents,
		llms.WithStreamingFunc(func(ctx context.Context, chunk []byte) error {
			fmt.Println("stream", string(chunk))
			return nil
		}))
	if err != nil {
		fmt.Println("绑定传入json数据信息异常", err)
		c.JSON(http.StatusBadRequest, gin.H{
			"msg": err.Error(),
		})
		return

	}
	str, _ := json.Marshal(content)
	fmt.Println("content", str)
	c.JSON(http.StatusOK, gin.H{
		"msg": content.Choices[0].Content,
	})
}
